//
// Created by ECNU on 2024/2/20.
//

#ifndef NANO_FAISS_INDEX_H
#define NANO_FAISS_INDEX_H

#include <nanofaiss/MetricType.h>
#include <cstdio>
#include <sstream>
#include <string>
#include <typeinfo>

namespace faiss {


struct Index {
    using component_t = float;
    using distance_t = float;

    int d; // dimension
    idx_t ntotal; // total number of indexed vectors

    bool is_trained; // whether is trained or need to be trained

    MetricType metric_type; // metric type is an enum
    [[maybe_unused]] float metric_arg; // metric arguments for lp

    explicit Index(idx_t d = 0, MetricType metric = MetricType::METRIC_L2)
            : d(d),
              ntotal(0),
              is_trained(true),
              metric_type(metric),
              metric_arg(0) {}

    virtual ~Index() = default;


    /** Perform training on a set of vectors
     * @param n     number of vectors
     * @param x     original vectors, is a n * d array
     */
    virtual void train(idx_t n, const float *x);


    /** Add n vector to the index
     *
     * Vectors are assigned id with ntotal ... ntotal + n - 1
     * @param n     number of vectors
     * @param x     original vectors, is a n * d array
     */
    virtual void add(idx_t n, const float *x) = 0;


    virtual void add_with_ids(idx_t n, const float *x, const idx_t *xids);


    /** n query vector x, search for k nearest
     * @param n             number of query
     * @param x             query vector
     * @param k             k nearest
     * @param distances     output distances, n * k
     * @param labels        output k nearest id, n * k
     */
    virtual void search(
            idx_t n,
            const float *x,
            idx_t k,
            float *distances,
            idx_t *labels) const = 0;


    /** same as search, but only return ids
     * @param n         number of query
     * @param x         original vectors
     * @param labels    output k nearest id
     * @param k         k nearest
     */
    virtual void assign(idx_t n, const float *x, idx_t *labels, idx_t k = 1) const;

    /// remove all the elements from the database
    virtual void reset() = 0;

    /// size of produced codes in bytes(in fact the vector is coded to code)
    virtual size_t sa_code_size() const;

    /** Encode a set of vector
     * @param n     number of vector
     * @param x     input vectors
     * @param bytes output encoded vectors n * sa_code_size
     */
    virtual void sa_encode(idx_t n, const float *x, uint8_t *bytes) const;

    /** decode a set of vectors
     * @param n     number of vector
     * @param bytes nput encoded vectors, size n * sa_code_size()
     * @param x     output vectors, size n * d
     */
    virtual void sa_decode(idx_t n, const uint8_t *bytes, float *x) const;
};


}

#endif // NANO_FAISS_INDEX_H
